import sys

try:
  import cv
  import Image
  import numpy
  from numpy import array
except ImportError, e:
  print >>sys.stderr, "You must install all dependencies: ", str(e)
  sys.exit(0)
else:
  print "Dependencies satisfied."

class Webcam:
  def __init__(self, camid=0, resolution='low'):
    self.cam = cv.CaptureFromCAM(camid)
    if resolution.startswith('lo'):
      width, height = 320, 240
    elif resolution == 'normal':
      width, height = 640, 480
    elif resolution.startswith('hi'):
      width, height = 1024, 960
    else:
      width, height = 640, 480
      
    cv.SetCaptureProperty(self.cam, cv.CV_CAP_PROP_FRAME_WIDTH, width)
    cv.SetCaptureProperty(self.cam, cv.CV_CAP_PROP_FRAME_HEIGHT, height)
    
  def __del__(self):
    del self.cam
    
  def stream(self, format='numpy'):
    while True:
      cv_image = cv.QueryFrame(self.cam)
      size = cv.GetSize(cv_image)
      (r, g, b) = [cv.CreateImage(size, cv_image.depth, 1) for i in range(3)]
      cv.Split(cv_image, b, g, r, None)
      cv.Merge(r, g, b, None, cv_image)

      pil_image = Image.fromstring('RGB', cv.GetSize(cv_image), cv_image.tostring())
      if format == 'pil':
        yield pil_image
      elif format == 'numpy':
        yield numpy.asarray(pil_image)
      else:
        raise Exception("Unknown format %s" % format)
  
  def __del__(self):
    del self.cam


def convert_pil_numpy(pil):
  return numpy.asarray(pil)

def convert_numpy_pil(numpy_image):
  (h,w) = numpy_image.shape[:2]
  return Image.fromarray(numpy_image)

def color_channel(numpy_image, color):
  x = numpy.array(numpy_image)
  if color == 'r':
    x[:,:,[1,2]] = 0
  elif color == 'g':
    x[:,:,[0,2]] = 0
  elif color == 'b':
    x[:,:,[0,1]] = 0
  return x

def apply_filter(numpy_image, filter, func=None):
  assert len(numpy_image.shape) == 3
  (h, w) = numpy_image.shape[:2]
  I = numpy_image.reshape((h*w, 3))
  X = numpy.dot(I, filter)
  if func:
    X = func(X)
  return X.reshape((h,w,3)).astype('uint8')

def make_filter(color, grayscale=False):
  if color == 'all':
    F = array([[1,0,0], [0, 1, 0], [0, 0, 1]])
  elif color == 'red':
    F = array([
      [1, 0, 0], 
      [0, 0, 0], 
      [0, 0, 0]])
  elif color == 'green':
    F = array([
      [0, 0, 0], 
      [0, 1, 0], 
      [0, 0, 0]])
  elif color == 'blue':
    F = array([
      [0, 0, 0], 
      [0, 0, 0], 
      [0, 0, 1]])
  if grayscale:
    F = numpy.dot(F, array([
        [0.33, 0.33, 0.33],
        [0.33, 0.33, 0.33],
        [0.34, 0.34, 0.34]]))
  return F
